#include <rclcpp/rclcpp.hpp>
#include <std_msgs/msg/float64.hpp>
#include <std_msgs/msg/float64_multi_array.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Vector3.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <string>

/*
 * ┌───────────────────────────────────────────────────────────────┐
 * │  force_td_servo_node（仅 TD + P 控制，用于快速验证接触平滑性） │
 * └───────────────────────────────────────────────────────────────┘
 * 功能：
 *   1) 订阅目标力 /fr_cmd（或用参数 force_td.target_force）
 *   2) TD（Tracking Differentiator）把 Fr 平滑成 v1（参考力）与 v2（其导数）
 *   3) 订阅 /wrench_comp 测得力，投影到接触法向，得到 F_meas
 *   4) P 控制：e = v1 - F_meas，vx = clamp(kp * e, -vmax, vmax)
 *   5) 发布速度到 servo_node/delta_twist_cmds（沿接触轴）
 *
 * 说明：
 *   - 这是“最小可用”的验证版：只引入 TD，其他（ESO/NLSEF/阻抗）暂不启用。
 *   - 目的：验证“加 TD 后，上力更平滑、冲击更小、P 控制更稳定”。
 *   - v1（N）：平滑参考力；v2（N/s）：参考力导数（这里仅记录观察，不参与控制）。
 */

static inline double clamp(double v, double lo, double hi) {
  return std::max(lo, std::min(hi, v));
}
static inline std::array<double,3> normalize3(std::array<double,3> v){
  double n = std::sqrt(v[0]*v[0]+v[1]*v[1]+v[2]*v[2]);
  if (n > 1e-9) { v[0]/=n; v[1]/=n; v[2]/=n; }
  return v;
}
static inline double dot3(const std::array<double,3>& a, const std::array<double,3>& b){
  return a[0]*b[0] + a[1]*b[1] + a[2]*b[2];
}

class ForceTDServoNode : public rclcpp::Node {
public:
  ForceTDServoNode() : Node("force_td_servo_node")
  {
    // ================= 参数 =================
    // 采样/控制周期
    dt_    = this->declare_parameter("sampling.dt", 0.01); // s

    // TD 参数（二阶线性跟踪器：临界阻尼默认 zeta=1）
    omega_ = this->declare_parameter("force_td.omega_td", 4.0);   // s^-1
    zeta_  = this->declare_parameter("force_td.zeta", 1.0);       // 1=临界阻尼
    target_force_param_ =
             this->declare_parameter("force_td.target_force", 0.8); // N
    init_from_first_cmd_ =
             this->declare_parameter("force_td.init_from_first_cmd", true);

    // 简单 P 控制与限速
    kp_    = this->declare_parameter("p_ctrl.kp", 0.10);   // (m/s)/N
    vmax_  = this->declare_parameter("p_ctrl.vmax", 0.05); // m/s

    // 坐标系与接触法向
    cmd_frame_  = this->declare_parameter<std::string>("frames.cmd_frame",  "base_link");
    axis_frame_ = this->declare_parameter<std::string>("frames.axis_frame", "base_link");
    contact_axis_ = {
      this->declare_parameter("frames.contact_axis_xyz[0]", 1.0),
      this->declare_parameter("frames.contact_axis_xyz[1]", 0.0),
      this->declare_parameter("frames.contact_axis_xyz[2]", 0.0)
    };
    contact_axis_ = normalize3(contact_axis_);

    // 话题
    twist_topic_  = this->declare_parameter<std::string>("topics.twist_out",  "servo_node/delta_twist_cmds");
    wrench_topic_ = this->declare_parameter<std::string>("topics.wrench_in",  "/wrench_comp");
    fr_cmd_topic_ = this->declare_parameter<std::string>("topics.fr_cmd_in",  "/fr_cmd");

    // 日志
    throttle_ms_  = this->declare_parameter("logging.throttle_ms", 200);
    publish_debug_= this->declare_parameter("logging.publish_debug", true);

    // ================= 打印启动参数 =================
    const double Trise_est = 2.3 / std::max(1e-6, omega_);
    RCLCPP_INFO(this->get_logger(),
      "force_td_servo_node 启动：\n"
      "  sampling.dt=%.4f s（控制周期）\n"
      "  force_td.omega_td=%.3f s^-1（TD 自然频率，越大过渡越快）\n"
      "  force_td.zeta=%.3f（TD 阻尼比，1=临界阻尼、无超调）\n"
      "  force_td.target_force=%.3f N（无 /fr_cmd 输入时使用）\n"
      "  force_td.init_from_first_cmd=%s（是否用首帧 /fr_cmd 初始化 v1）\n"
      "  p_ctrl.kp=%.3f (m/s)/N（P 控制比例系数）\n"
      "  p_ctrl.vmax=%.3f m/s（速度限幅）\n"
      "  frames.cmd_frame=%s（发布 Twist 的坐标系）\n"
      "  frames.axis_frame=%s（接触法向定义的坐标系）\n"
      "  frames.contact_axis_xyz=[%.3f, %.3f, %.3f]（接触法向，单位向量）\n"
      "  topics.twist_out=%s\n"
      "  topics.wrench_in=%s\n"
      "  topics.fr_cmd_in=%s\n"
      "  推算 TD 抬升时间（10%%→90%%）≈ %.3f s",
      dt_, omega_, zeta_, target_force_param_,
      init_from_first_cmd_?"true":"false",
      kp_, vmax_, cmd_frame_.c_str(), axis_frame_.c_str(),
      contact_axis_[0], contact_axis_[1], contact_axis_[2],
      twist_topic_.c_str(), wrench_topic_.c_str(), fr_cmd_topic_.c_str(),
      Trise_est);

    // ================= I/O =================
    twist_pub_ = this->create_publisher<geometry_msgs::msg::TwistStamped>(twist_topic_, 10);
    debug_pub_ = this->create_publisher<std_msgs::msg::Float64MultiArray>("/fr_td/debug", 10);

    // 力：SensorDataQoS，必要时做 TF 到 cmd_frame 并投影到接触轴
    wrench_sub_ = this->create_subscription<geometry_msgs::msg::WrenchStamped>(
      wrench_topic_, rclcpp::SensorDataQoS(),
      [this](geometry_msgs::msg::WrenchStamped::SharedPtr msg){ onWrench(msg); });

    // 目标力
    fr_cmd_sub_ = this->create_subscription<std_msgs::msg::Float64>(
      fr_cmd_topic_, 10,
      [this](std_msgs::msg::Float64::SharedPtr m){
        last_fr_cmd_ = m->data; has_fr_cmd_ = true;
        if (!td_inited_ && init_from_first_cmd_) {
          v1_ = last_fr_cmd_; v2_ = 0.0; td_inited_ = true;
          RCLCPP_INFO(this->get_logger(),
            "TD 初始化：使用首次 /fr_cmd=%.3f N 作为 v1 初值（init_from_first_cmd=true）", last_fr_cmd_);
        }
      });

    // TF
    tf_buffer_   = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

    // TD 初始（若不等待首帧，则用参数初始化）
    if (!init_from_first_cmd_) { v1_ = target_force_param_; v2_ = 0.0; td_inited_ = true; }

    // 控制定时器
    timer_ = this->create_wall_timer(
      std::chrono::duration<double>(dt_), std::bind(&ForceTDServoNode::controlTick, this));
  }

private:
  // ---------- 力回调：变换到 cmd_frame 并投影到接触法向 ----------
  void onWrench(const geometry_msgs::msg::WrenchStamped::SharedPtr& msg)
  {
    geometry_msgs::msg::WrenchStamped w = *msg;
    // 若力的 frame 与 cmd_frame 不同，做 TF 变换
    if (msg->header.frame_id != cmd_frame_) {
      try {
        auto T = tf_buffer_->lookupTransform(cmd_frame_, msg->header.frame_id, tf2::TimePointZero);
        tf2::doTransform(*msg, w, T);
      } catch (const tf2::TransformException& ex) {
        RCLCPP_WARN_THROTTLE(this->get_logger(), *this->get_clock(), 2000,
          "TF 变换测力到 cmd_frame 失败：%s（将使用原始 frame）", ex.what());
        w = *msg; // 失败就用原数据（假定误差可接受）
      }
    }
    // 接触法向在 cmd_frame 下的单位向量
    std::array<double,3> n_cmd = contactAxisInCmd();
    // 取力向量并投影
    std::array<double,3> F = { w.wrench.force.x, w.wrench.force.y, w.wrench.force.z };
    last_F_meas_scalar_ = dot3(n_cmd, F);
    have_force_ = true;
  }

  // ---------- TD 连续模型的欧拉离散 ----------
  // v1_dot = v2
  // v2_dot = -2*zeta*omega*v2 - omega^2*(v1 - Fr)
  inline void td_step(double Fr, double dt, double omega, double zeta,
                      double& v1, double& v2)
  {
    const double v1_dot = v2;
    const double v2_dot = -2.0*zeta*omega*v2 - (omega*omega)*(v1 - Fr);
    v1 += dt * v1_dot;
    v2 += dt * v2_dot;
  }

  // ---------- 主控制循环 ----------
  void controlTick()
  {
    // 1) 目标力：若无 /fr_cmd 则用参数
    const double Fr = has_fr_cmd_ ? last_fr_cmd_ : target_force_param_;

    // 2) TD 初始化（若选择等待首帧且尚未初始化）
    if (!td_inited_) {
      v1_ = Fr; v2_ = 0.0; td_inited_ = true;
      RCLCPP_INFO(this->get_logger(),
        "TD 初始化：使用当前 Fr=%.3f N 作为 v1 初值（未收到 /fr_cmd，使用参数）", Fr);
    }

    // 3) TD 更新：得到 v1（平滑参考力）, v2（参考力导）
    td_step(Fr, dt_, omega_, zeta_, v1_, v2_);

    // 4) 若还没收到力，就不动，等力有效
    if (!have_force_) {
      publishTwist(0.0);
      if (publish_debug_) publishDebug(Fr, /*F_meas*/NAN, v1_, v2_, /*err*/NAN, /*vx*/0.0);
      RCLCPP_INFO_THROTTLE(this->get_logger(), *this->get_clock(), throttle_ms_,
        "[等待测力] 尚未收到 /wrench_comp；仅输出 0 速度。Fr=%.3f N | v1=%.3f N | v2=%.3f N/s",
        Fr, v1_, v2_);
      return;
    }

    // 5) P 控制：跟踪 v1
    const double Fm  = last_F_meas_scalar_;
    const double err = v1_ - Fm;                 // 力误差（参考力 - 实测力）
    double v_cmd     = clamp(kp_ * err, -vmax_, vmax_); // 限速

    // 6) 下发速度
    publishTwist(v_cmd);

    // 7) 调试发布与节流日志（字段名+单位都写清楚）
    if (publish_debug_) publishDebug(Fr, Fm, v1_, v2_, err, v_cmd);

    RCLCPP_INFO_THROTTLE(this->get_logger(), *this->get_clock(), throttle_ms_,
      "[TD+P] Fr_raw=%.3f N（目标力输入） | v1=%.3f N（TD平滑参考力） | v2=%.3f N/s（TD参考力导） | "
      "F_meas=%.3f N（实测力-投影到接触轴） | err=%.3f N（力误差） | "
      "kp=%.3f (m/s)/N（比例系数） | vmax=%.3f m/s（限速） | v_cmd=%.3f m/s（沿接触轴的速度命令） | "
      "omega_td=%.3f s^-1（TD自然频率） | zeta=%.3f（TD阻尼比） | dt=%.4f s（控制周期）",
      Fr, v1_, v2_, Fm, err, kp_, vmax_, v_cmd, omega_, zeta_, dt_);
  }

  // ---------- 工具：获取接触轴在 cmd_frame 下的单位向量 ----------
  std::array<double,3> contactAxisInCmd()
  {
    if (axis_frame_ == cmd_frame_) return contact_axis_;
    try {
      auto T = tf_buffer_->lookupTransform(cmd_frame_, axis_frame_, tf2::TimePointZero);
      tf2::Quaternion q; tf2::fromMsg(T.transform.rotation, q);
      tf2::Vector3 v_axis(contact_axis_[0], contact_axis_[1], contact_axis_[2]);
      tf2::Vector3 v_cmd = tf2::quatRotate(q, v_axis);
      std::array<double,3> n{ v_cmd.x(), v_cmd.y(), v_cmd.z() };
      return normalize3(n);
    } catch (...) {
      return contact_axis_; // TF 失败时退回原轴
    }
  }

  // ---------- 发布 Twist（沿接触轴的线速度） ----------
  void publishTwist(double v_axis_cmd)
  {
    auto n_cmd = contactAxisInCmd();
    geometry_msgs::msg::TwistStamped out;
    out.header.stamp = now();
    out.header.frame_id = cmd_frame_;
    out.twist.linear.x = n_cmd[0] * v_axis_cmd;
    out.twist.linear.y = n_cmd[1] * v_axis_cmd;
    out.twist.linear.z = n_cmd[2] * v_axis_cmd;
    out.twist.angular.x = 0.0;
    out.twist.angular.y = 0.0;
    out.twist.angular.z = 0.0;
    twist_pub_->publish(out);
  }

  // ---------- 发布 debug 数组 ----------
  void publishDebug(double Fr, double Fm, double v1, double v2, double err, double vcmd)
  {
    std_msgs::msg::Float64MultiArray m;
    // [0]Fr_raw(N), [1]F_meas(N), [2]v1_ref(N), [3]v2_ref(N/s),
    // [4]err(N), [5]v_cmd(m/s), [6]kp((m/s)/N), [7]vmax(m/s),
    // [8]omega_td(s^-1), [9]zeta, [10]dt(s)
    m.data = { Fr, Fm, v1, v2, err, vcmd, kp_, vmax_, omega_, zeta_, dt_ };
    debug_pub_->publish(m);
  }

private:
  // 参数
  double dt_{0.01};
  double omega_{4.0}, zeta_{1.0};
  double target_force_param_{0.8};
  bool   init_from_first_cmd_{true};
  double kp_{0.10}, vmax_{0.05};
  std::string cmd_frame_{"base_link"}, axis_frame_{"base_link"};
  std::array<double,3> contact_axis_{1.0,0.0,0.0};
  std::string twist_topic_, wrench_topic_, fr_cmd_topic_;
  int throttle_ms_{200};
  bool publish_debug_{true};

  // IO
  rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_pub_;
  rclcpp::Publisher<std_msgs::msg::Float64MultiArray>::SharedPtr debug_pub_;
  rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr wrench_sub_;
  rclcpp::Subscription<std_msgs::msg::Float64>::SharedPtr fr_cmd_sub_;

  // TF
  std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;

  // TD 状态
  bool   td_inited_{false};
  double v1_{0.0}, v2_{0.0}; // v1:参考力(N), v2:参考力导(N/s)

  // 输入缓存
  bool   has_fr_cmd_{false};
  double last_fr_cmd_{0.8};
  bool   have_force_{false};
  double last_F_meas_scalar_{0.0}; // 投影后的实测力(N)

  // 定时器
  rclcpp::TimerBase::SharedPtr timer_;
};

int main(int argc, char** argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ForceTDServoNode>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
